from os import listdir
from os.path import join
import os
import re
import base64
import pandas as pd
import imghdr

class DecodeBase64(object):
	
	def __init__(self):
		self.idno = []
	def decode(self, path1, output_path):
		pattern = re.compile(r'[\d]{18}([X\d]{18}|[X\d]{15}).dat')
		idno_all = []
		count = 0
		no_type = 0
		for f in listdir(path1):
			filepath = join(path1, f)
			m = pattern.search(f)
			try:
				idno = m.group(1)
			except:
				print('IDNO ERROR:' + f)
				continue
			read_file = open(filepath, 'rb')
			lines = read_file.readlines()
			word = lines[0].strip()
			try:
				b = base64.decodestring(word)

				#print(b[:2])

				if b[:2] == b'\xff\xd8':
					outputPath = output_path+idno+'.jpeg'
				elif b[:2] == b'8950':
					outputPath = output_path+idno+'.png'
				elif b[:2] == b'4749':
					outputPath = output_path+idno+'.gif'
				elif b[:2] == b'424d':
					outputPath = output_path+idno+'bmp'
				elif b[:2] == b'd8e0':
					outputPath = output_path+idno+'jfif'
				else:
					outputPath = output_path+idno+'.jpeg'
					print(b[:2], filepath)
					print('no type')
					no_type += 1
			except Exception as e:
				print ("BASE64 ERROR:" + f)
				print (e)
				continue
		
			write_file = open(outputPath, 'wb')
			write_file.write(b)
			write_file.close()
			idno_all.append(idno)
			count += 1
		
		print(no_type)
		self.idno = idno_all
		#self._outputData(self.idno, 'idno.csv')

	def _outputData(self, data, path1):
		df = pd.DataFrame(data)
		df.to_csv(path1, index = False, sep = ',', header = False)

if __name__ == '__main__':
	path1 = '../smoke_data/data/smoke/'
	output_path = '../smoke_data/smoke_output/smoke_output/'
	decoder = DecodeBase64()
	decoder.decode(path1, output_path)
